import torch
from torch.utils.data import DataLoader
import torchmetrics  # conda install torchmetrics
from net.DLinkNet import DLinkNet34
import Dataset
import tqdm


class Metrics():
    def __init__(self, task="binary", num_classes=2):
        self.task = task
        self.num_classes = num_classes
        self.ConfusionMatrix = torchmetrics.ConfusionMatrix(task=self.task, num_classes=self.num_classes)

    def OA_metrics(self, TN, FP, FN, TP):
        return (TP + TN) / (TP + FP + TN + FN)

    def P_metrics(self, TP, FP):
        return TP / (TP + FP)

    def R_metrics(self, TP, FN):
        return TP / (TP + FN)

    def F1_metrics(self, P, R):
        return (2 * P * R) / (P + R)

    def IoU_metrics(self, TP, FP, FN):
        return TP / (TP + FP + FN)

    def metrics(self, y_pred, y_true):
        threshold = 0.5  # You can adjust this threshold based on your requirements
        y_pred_bin = (y_pred > threshold).int()
        y_true_bin = (y_true > threshold).int()

        TN, FP, FN, TP = self.ConfusionMatrix(y_pred_bin, y_true_bin).ravel()
        OA = self.OA_metrics(TN, FP, FN, TP)
        P = self.P_metrics(TP, FP)
        R = self.R_metrics(TP, FN)
        F1 = self.F1_metrics(P, R)
        IoU = self.IoU_metrics(TP, FP, FN)

        return P, R, OA, F1, IoU


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    batch_size = 1
    num_worker = 0

    total_P = 0
    total_R = 0
    total_OA = 0
    total_F1 = 0
    total_IoU = 0

    test_dataset = Dataset.SegmentationDataset(where='test', seq=None)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_worker)
    lens = test_dataset.__len__()

    model = DLinkNet34().to(device)
    model.load_state_dict(torch.load('./model/class12.pt', map_location=device))
    model.eval()

    model_Metrics = Metrics(task="binary", num_classes=2)
    for i, (x, y, name) in enumerate(tqdm.tqdm(test_loader)):
        # 获取图像
        x = x.to(device, dtype=torch.float32)
        y = torch.squeeze(y, dim=0)
        y = torch.squeeze(y, dim=0)

        # 获取预测值
        y_pred = model(x)
        y_pred = torch.squeeze(y_pred, dim=0)
        y_pred = torch.squeeze(y_pred, dim=0).cpu()

        P, R, OA, F1, IoU = model_Metrics.metrics(y_pred, y)
        if torch.isnan(torch.stack([P, R, OA, F1, IoU])).any():
            lens = lens - 1
            continue
        else:
            total_P += P
            total_R += R
            total_OA += OA
            total_F1 += F1
            total_IoU += IoU

        # print(P, R, OA, F1, IoU, file=file)

avg_P = total_P / lens
avg_R = total_R / lens
avg_OA = total_OA / lens
avg_F1 = total_F1 / lens
avg_IoU = total_IoU / lens

print('avg_P:', avg_P)
print('avg_R:', avg_R)
print('avg_OA:', avg_OA)
print('avg_F1:', avg_F1)
print('avg_IoU:', avg_IoU)
